from fastapi import FastAPI, HTTPException, Depends, Header, UploadFile, File, Form, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional
import json
from jose import jwt
from datetime import datetime, timedelta
import os
import shutil
import zipfile
from io import BytesIO
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import traceback
import random

app = FastAPI()

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")

# 添加根路由，返回前端页面
@app.get("/")
async def read_root():
    return FileResponse("static/index.html")

# 模拟的物体类别
MOCK_CLASSES = ['人', '车', '狗', '猫', '椅子', '桌子', '电脑', '手机']

# 初始化人脸检测器
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# 基础目录配置
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
UPLOAD_ROOT = os.path.join(BASE_DIR, "uploads")
FACE_DB_ROOT = os.path.join(UPLOAD_ROOT, "face_db")  # 人脸数据库根目录

# 为不同类型的图片创建不同的上传目录
UPLOAD_DIRS = {
    "category": os.path.join(UPLOAD_ROOT, "categories"),
    "persons": os.path.join(UPLOAD_ROOT, "persons"),
    "face": os.path.join(UPLOAD_ROOT, "faces"),
    "objects": os.path.join(UPLOAD_ROOT, "objects")  # 添加物体检测目录
}

# 创建必要的目录
for dir_path in UPLOAD_DIRS.values():
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

# 确保人脸数据库目录存在
if not os.path.exists(FACE_DB_ROOT):
    os.makedirs(FACE_DB_ROOT)

# 配置静态文件服务
app.mount("/uploads", StaticFiles(directory=UPLOAD_ROOT), name="uploads")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 数据模型
class UserLogin(BaseModel):
    username: str
    password: str

class Category(BaseModel):
    name: str
    desc: Optional[str] = None

class CategoryResponse(Category):
    id: int
    created_at: str
    created_by: str

# JSON文件路径
CATEGORIES_FILE = 'categories.json'
FACE_NAMES_FILE = 'face_names.json'

# 图片模型
class ImageInfo(BaseModel):
    id: int
    filename: str
    category_id: int
    uploaded_by: str
    uploaded_at: str

# 人脸特征数据存储
FACE_FEATURES_FILE = 'face_features.json'

# 人脸数据存储
FACES_FILE = 'faces.json'

# 更新图片分类的请求模型
class UpdateCategoryRequest(BaseModel):
    category_id: Optional[int] = None
    type: str = "category"

# 行人检测的请求模型
class PersonDetectionResponse(BaseModel):
    id: int
    filename: str
    annotated_image: str
    persons: List[dict]
    uploaded_by: str
    uploaded_at: str

# JWT 配置
SECRET_KEY = "your-secret-key"  # 在生产环境中应该使用更安全的密钥
ALGORITHM = "HS256"

def create_token(data: dict):
    to_encode = data.copy()
    token = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return token

def verify_token(token: str = Header(...)):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return payload.get("username")
    except Exception:
        raise HTTPException(
            status_code=401,
            detail="Invalid authentication credentials"
        )

def create_face_folder(name):
    """创建人名对应的文件夹"""
    folder_path = os.path.join(FACE_DB_ROOT, name)
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    return folder_path

def get_face_encodings(image_path):
    """获取图片中的人脸特征"""
    # 使用 numpy 和 imdecode 处理图片，支持中文路径
    image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    if image is None:
        raise Exception("无法读取图片")
    
    # 转换为RGB格式（OpenCV默认是BGR）
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 检测人脸
    faces = face_cascade.detectMultiScale(rgb_image, 1.1, 4)
    
    # 提取每个人脸的特征
    face_features = []
    for (x, y, w, h) in faces:
        face_roi = rgb_image[y:y+h, x:x+w]
        # 这里可以添加更多的特征提取方法
        # 简单起见，我们只保存人脸区域的平均颜色作为特征
        feature = [float(x) for x in cv2.mean(face_roi)[:3]]  # 取RGB均值
        face_features.append({
            'feature': feature,
            'bbox': {
                'x': int(x),
                'y': int(y),
                'width': int(w),
                'height': int(h)
            }
        })
    
    return face_features

def compare_faces(face_feature, db_features):
    """比较人脸特征，返回最匹配的结果"""
    best_match = None
    best_score = float('inf')
    
    for db_face in db_features:
        # 计算特征向量的欧氏距离
        # 确保使用 Python 原生类型进行计算
        diff = sum((float(a) - float(b)) ** 2 for a, b in zip(face_feature, db_face['feature']))
        if diff < best_score:
            best_score = diff
            best_match = db_face
    
    return best_match, best_score

# 初始化分类数据
def init_categories():
    if not os.path.exists(CATEGORIES_FILE):
        default_categories = [
            {
                "id": 1,
                "name": "全部图片",
                "desc": "所有上传的图片",
                "created_at": datetime.now().isoformat(),
                "created_by": "admin"
            },
            {
                "id": 2,
                "name": "风景",
                "desc": "风景照片",
                "created_at": datetime.now().isoformat(),
                "created_by": "admin"
            },
            {
                "id": 3,
                "name": "人物",
                "desc": "人物照片",
                "created_at": datetime.now().isoformat(),
                "created_by": "admin"
            },
            {
                "id": 4,
                "name": "建筑",
                "desc": "建筑照片",
                "created_at": datetime.now().isoformat(),
                "created_by": "admin"
            }
        ]
        save_categories(default_categories)

# 保存分类到JSON文件
def save_categories(categories):
    with open(CATEGORIES_FILE, 'w', encoding='utf-8') as f:
        json.dump(categories, f, ensure_ascii=False, indent=2)

# 读取分类列表
def load_categories():
    if os.path.exists(CATEGORIES_FILE):
        with open(CATEGORIES_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    return []

# 读取图片数据
def load_images(type):
    """加载指定类型的图片数据"""
    json_file = f'images_{type}.json'
    if os.path.exists(json_file):
        with open(json_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    return []

# 保存图片数据
def save_images(images, type):
    """保存指定类型的图片数据"""
    json_file = f'images_{type}.json'
    with open(json_file, 'w', encoding='utf-8') as f:
        json.dump(images, f, ensure_ascii=False, indent=2)

# API路由
@app.get("/api/categories", response_model=List[CategoryResponse])
async def get_categories(username: str = Depends(verify_token)):
    return load_categories()

@app.post("/api/categories", response_model=CategoryResponse)
async def create_category(category: Category, username: str = Depends(verify_token)):
    categories = load_categories()
    new_id = max([c["id"] for c in categories], default=0) + 1
    
    new_category = {
        "id": new_id,
        "name": category.name,
        "desc": category.desc,
        "created_at": datetime.now().isoformat(),
        "created_by": username
    }
    
    categories.append(new_category)
    save_categories(categories)
    return new_category

@app.delete("/api/categories/{category_id}")
async def delete_category(category_id: int, username: str = Depends(verify_token)):
    categories = load_categories()
    
    # 找到要删除的分类
    category = next((c for c in categories if c["id"] == category_id), None)
    if not category:
        raise HTTPException(status_code=404, detail="分类不存在")
    
    # 检查是否为创建者
    if category["created_by"] != username:
        raise HTTPException(status_code=403, detail="没有权限删除此分类")
    
    # 删除分类
    categories = [c for c in categories if c["id"] != category_id]
    save_categories(categories)
    return {"message": "分类已删除"}

# 修改登录接口
@app.post("/api/login")
async def login(user: UserLogin):
    # 简单的用户验证，接受默认用户名密码: admin/admin
    if user.username == "admin" and user.password == "admin":
        # 创建token
        token = create_token({
            "username": user.username,
            "exp": datetime.utcnow() + timedelta(days=1)
        })
        return {
            "token": token,
            "username": user.username,
            "userId": 1
        }
    raise HTTPException(status_code=401, detail="用户名或密码错误")

# 初始化 HOG 检测器
hog = cv2.HOGDescriptor()
hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())

def cv2_img_add_text(img, text, left, top, text_color=(255, 255, 255), text_size=20):
    """使用PIL添加中文文字到图片上"""
    if isinstance(img, np.ndarray):
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img)
    
    # 加载中文字体文件（需要系统安装了相应的字体）
    fontStyle = ImageFont.truetype("SimHei.ttf", text_size, encoding="utf-8")
    draw.text((left, top), text, text_color, font=fontStyle)
    
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

def analyze_person(img, bbox):
    """分析行人的姿态、服装和动作"""
    x, y, w, h = bbox
    print(f"开始分析行人特征: bbox={bbox}")  # 调试信息
    
    try:
        # 确保坐标在图片范围内
        x = max(0, x)
        y = max(0, y)
        w = min(w, img.shape[1] - x)
        h = min(h, img.shape[0] - y)
        
        person_img = img[y:y+h, x:x+w]
        if person_img.size == 0:
            print("无法获取人物图像区域")  # 调试信息
            return {"pose": "未知", "clothing": "未知", "action": "未知"}

        print(f"人物图像区域大小: {person_img.shape}")  # 调试信息

        # 姿态分析
        aspect_ratio = h / w
        print(f"宽高比: {aspect_ratio}")  # 调试信息
        
        if aspect_ratio > 2.2:
            pose = "站立"
        elif aspect_ratio < 1.5:
            pose = "坐着"
        else:
            pose = "站立"
        print(f"判断姿态: {pose}")  # 调试信息
        
        # 服装颜色分析
        hsv_img = cv2.cvtColor(person_img, cv2.COLOR_BGR2HSV)
        upper_body = hsv_img[:h//2, :]  # 分析上半身
        
        # 定义HSV颜色范围
        color_ranges = {
            "红色": ([0, 50, 50], [10, 255, 255]),  # 放宽红色的范围
            "蓝色": ([100, 50, 50], [130, 255, 255]),  # 放宽蓝色的范围
            "绿色": ([40, 50, 50], [80, 255, 255]),
            "白色": ([0, 0, 180], [180, 30, 255]),  # 调整白色的范围
            "黑色": ([0, 0, 0], [180, 255, 50]),
            "棕色": ([10, 50, 20], [20, 255, 200])
        }
        
        max_pixels = 0
        clothing = "未知"
        total_pixels = upper_body.size / 3  # 总像素数
        
        print(f"上半身区域总像素数: {total_pixels}")  # 调试信息
        
        for color_name, (lower, upper) in color_ranges.items():
            lower = np.array(lower)
            upper = np.array(upper)
            mask = cv2.inRange(upper_body, lower, upper)
            pixel_count = np.sum(mask > 0)
            pixel_ratio = pixel_count / total_pixels
            print(f"颜色 {color_name} 的像素比例: {pixel_ratio:.2%}")  # 调试信息
            
            # 降低阈值到5%
            if pixel_count > max_pixels and pixel_ratio > 0.05:  # 只需要5%的像素
                max_pixels = pixel_count
                clothing = f"{color_name}衣服"
        
        print(f"判断服装: {clothing}")  # 调试信息
        
        # 动作分析
        try:
            if w > h * 0.4:
                action = "站立"
            else:
                action = "行走"
            print(f"判断动作: {action}")  # 调试信息
            
        except Exception as e:
            print(f"动作分析错误: {str(e)}")  # 调试信息
            action = "未知"
        
        result = {
            "pose": pose,
            "clothing": clothing,
            "action": action
        }
        print(f"分析结果: {result}")  # 调试信息
        return result
        
    except Exception as e:
        print(f"特征分析错误: {str(e)}")
        traceback.print_exc()  # 打印完整的错误堆栈
        return {
            "pose": "未知",
            "clothing": "未知",
            "action": "未知"
        }

def save_annotation_txt(image_path, persons, image_size):
    """保存YOLO格式的标注文件"""
    try:
        # 将图片路径的扩展名改为.txt
        txt_path = os.path.splitext(image_path)[0] + '.txt'
        
        # 获取图片尺寸
        img_width, img_height = image_size
        
        with open(txt_path, 'w', encoding='utf-8') as f:
            for person in persons:
                # 转换为YOLO格式（归一化坐标）
                x = person["bbox"]["x"] / 100 * img_width
                y = person["bbox"]["y"] / 100 * img_height
                w = person["bbox"]["width"] / 100 * img_width
                h = person["bbox"]["height"] / 100 * img_height
                
                # 计算中心点坐标和宽高（归一化到0-1）
                x_center = (x + w/2) / img_width
                y_center = (y + h/2) / img_height
                width = w / img_width
                height = h / img_height
                
                # 限制在0-1范围内
                x_center = max(0, min(1, x_center))
                y_center = max(0, min(1, y_center))
                width = max(0, min(1, width))
                height = max(0, min(1, height))
                
                # 写入YOLO格式：class x_center y_center width height
                # 这里class=0表示"person"类
                line = f"0 {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
                
                # 添加额外的属性信息（作为注释）
                attributes = f"# confidence:{person['confidence']:.2f} pose:{person['pose']} clothing:{person['clothing']} action:{person['action']}"
                
                f.write(f"{line} {attributes}\n")
        
        print(f"保存标注文件: {txt_path}")
        
        # 同时生成一个classes.txt文件（如果不存在）
        classes_path = os.path.join(os.path.dirname(txt_path), 'classes.txt')
        if not os.path.exists(classes_path):
            with open(classes_path, 'w', encoding='utf-8') as f:
                f.write('person\n')
            print(f"创建类别文件: {classes_path}")
            
    except Exception as e:
        print(f"保存标注文件失败: {str(e)}")
        traceback.print_exc()

def detect_persons(image_path):
    print(f"开始处理图片: {image_path}")
    
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        raise Exception("无法读取图片")

    # 获取原始尺寸
    original_height, original_width = img.shape[:2]
    
    # 调整图片大小以提高检测效果
    height = 800
    scale = height / original_height
    width = int(original_width * scale)
    resized_img = cv2.resize(img, (width, height))

    # 行人检测
    boxes, weights = hog.detectMultiScale(
        resized_img, 
        winStride=(4, 4),
        padding=(8, 8), 
        scale=1.05,
        hitThreshold=0.3
    )

    # 创建标注图片的副本
    annotated_img = img.copy()
    
    # 计算缩放比例
    scale_x = original_width / width
    scale_y = original_height / height
    
    persons = []

    try:
        for i, (x, y, w, h) in enumerate(boxes):
            confidence = min(float(weights[i]), 1.0)
            
            if confidence > 0.1:
                # 将检测框坐标转换回原始尺寸
                orig_x = int(x * scale_x)
                orig_y = int(y * scale_y)
                orig_w = int(w * scale_x)
                orig_h = int(h * scale_y)

                # 计算相对坐标（百分比）
                rel_x = (orig_x / original_width) * 100
                rel_y = (orig_y / original_height) * 100
                rel_w = (orig_w / original_width) * 100
                rel_h = (orig_h / original_height) * 100

                # 绘制检测框
                cv2.rectangle(annotated_img, 
                            (orig_x, orig_y), 
                            (orig_x + orig_w, orig_y + orig_h), 
                            (0, 255, 0),  # BGR 绿色
                            2)  # 线条粗细

                # 分析行人特征
                analysis = analyze_person(img, (orig_x, orig_y, orig_w, orig_h))

                # 记录检测结果（使用相对坐标）
                person = {
                    "id": i + 1,
                    "confidence": float(confidence),
                    "bbox": {
                        "x": rel_x,
                        "y": rel_y,
                        "width": rel_w,
                        "height": rel_h
                    },
                    "pose": analysis["pose"],
                    "clothing": analysis["clothing"],
                    "action": analysis["action"]
                }
                persons.append(person)
                print(f"检测到行人 {i+1}:", person)

        # 保存标注后的图片
        annotated_filename = f"annotated_{os.path.basename(image_path)}"
        annotated_path = os.path.join(UPLOAD_DIRS["persons"], annotated_filename)
        cv2.imwrite(annotated_path, annotated_img)

        result = {
            "persons": persons,
            "annotated_image": f"persons/{annotated_filename}"
        }
        print("返回结果:", result)  # 调试信息
        return result

    except Exception as e:
        print(f"处理图片时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 人脸检测函数
def detect_faces(image_path):
    try:
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise Exception("无法读取图片")
        
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 检测人脸
        faces = face_cascade.detectMultiScale(gray, 1.1, 4)
        
        result = []
        for i, (x, y, w, h) in enumerate(faces):
            face_id = str(i + 1)
            
            # 在图片上画框
            cv2.rectangle(image, (x, y), (x+w, y+h), (255, 0, 0), 2)
            cv2.putText(image, f'Face {face_id}', (x, y-10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
            
            result.append({
                'id': face_id,
                'name': f'未知人脸 {face_id}',
                'name_id': None,
                'confidence': 0.8,
                'bbox': {
                    'x': int(x),
                    'y': int(y),
                    'width': int(w),
                    'height': int(h)
                }
            })
        
        # 保存标注图片
        annotated_filename = f"annotated_{os.path.basename(image_path)}"
        annotated_path = os.path.join(UPLOAD_DIRS['face'], annotated_filename)
        cv2.imwrite(annotated_path, image)
        
        return {
            'faces': result,
            'annotated_image': f"faces/{annotated_filename}"
        }
    except Exception as e:
        print(f"人脸检测失败: {str(e)}")
        traceback.print_exc()
        raise e

# 处理图片上传
@app.post("/api/upload")
async def upload_image(
    file: UploadFile = File(...),
    type: str = Form(...),
    category_id: Optional[int] = Form(None),
    username: str = Depends(verify_token)
):
    try:
        # 验证上传类型
        if type not in UPLOAD_DIRS:
            raise HTTPException(status_code=400, detail="无效的上传类型")
        
        # 验证文件类型
        if not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="只能上传图片文件")
        
        # 生成文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        random_str = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=8))
        filename = f"{timestamp}_{random_str}{os.path.splitext(file.filename)[1]}"
        
        # 确保上传目录存在
        upload_dir = UPLOAD_DIRS[type]
        if not os.path.exists(upload_dir):
            os.makedirs(upload_dir)
        
        # 保存文件
        file_path = os.path.join(upload_dir, filename)
        with open(file_path, "wb") as f:
            shutil.copyfileobj(file.file, f)
        
        # 读取现有图片数据
        images = load_images(type)
        
        # 创建新的图片记录
        new_image = {
            "id": len(images) + 1,
            "filename": filename,
            "category_id": category_id,
            "uploaded_by": username,
            "uploaded_at": datetime.now().isoformat()
        }
        
        # 添加新图片并保存
        images.append(new_image)
        save_images(images, type)
        
        # 返回完整的图片信息
        return {
            "id": new_image["id"],
            "filename": f"{type}/{filename}",
            "category_id": new_image["category_id"],
            "uploaded_by": new_image["uploaded_by"],
            "uploaded_at": new_image["uploaded_at"]
        }
    except Exception as e:
        print(f"上传失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 获取图片列表
@app.get("/api/images")
async def get_images(
    type: str = Query(default="category", description="图片类型"),
    category_id: Optional[int] = None,
    username: str = Depends(verify_token)
):
    try:
        # 修正类型名称
        if type == "person":
            type = "persons"
        
        # 验证类型
        if type not in UPLOAD_DIRS:
            raise HTTPException(status_code=400, detail="无效的图片类型")
        
        # 读取图片数据
        images = load_images(type)
        if images is None:
            images = []
        
        # 如果是分类类型且指定了分类ID，进行过滤
        if type == "category" and category_id is not None:
            images = [img for img in images if img.get("category_id") == category_id]
        
        # 确保文件路径正确
        for image in images:
            if "filename" in image and not image["filename"].startswith(type):
                image["filename"] = f"{type}/{os.path.basename(image['filename'])}"
            if "annotated_image" in image and not image["annotated_image"].startswith(type):
                image["annotated_image"] = f"{type}/{os.path.basename(image['annotated_image'])}"
        
        return images
        
    except Exception as e:
        print(f"获取图片列表失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 更新图片分类
@app.put("/api/images/{image_id}/category")
async def update_image_category(
    image_id: int,
    update_data: UpdateCategoryRequest,
    username: str = Depends(verify_token)
):
    try:
        # 验证图片类型
        if update_data.type not in UPLOAD_DIRS:
            raise HTTPException(status_code=400, detail="无效的图片类型")
        
        # 读取图片数据
        images = load_images(update_data.type)
        
        # 查找要更新的图片
        image = next((img for img in images if img["id"] == image_id), None)
        if not image:
            raise HTTPException(status_code=404, detail="图片不存在")
        
        # 更新分类
        image["category_id"] = update_data.category_id
        
        # 保存更新
        save_images(images, update_data.type)
        
        return image
    except Exception as e:
        print(f"更新分类失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 删除图片
@app.delete("/api/images/{image_id}")
async def delete_image(
    image_id: int,
    type: str = Query(..., description="图片类型"),
    username: str = Depends(verify_token)
):
    try:
        if type not in UPLOAD_DIRS:
            raise HTTPException(status_code=400, detail="无效的图片类型")
        
        images = load_images(type)
        image = next((img for img in images if img["id"] == image_id), None)
        
        if not image:
            raise HTTPException(status_code=404, detail="图片不存在")
            
        # 删除文件
        file_path = os.path.join(UPLOAD_DIRS[type], image["filename"])
        if os.path.exists(file_path):
            os.remove(file_path)
            
        # 更新数据
        images = [img for img in images if img["id"] != image_id]
        save_images(images, type)
        
        return {"message": "删除成功"}
    except Exception as e:
        print(f"删除失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 启动时初始化分类数据
@app.on_event("startup")
async def startup_event():
    init_categories()
    init_face_names()

# 加载人脸特征
def load_face_features():
    if os.path.exists(FACE_FEATURES_FILE):
        with open(FACE_FEATURES_FILE, 'r') as f:
            return json.load(f)
    return {}

# 保存人脸特征
def save_face_features(features):
    with open(FACE_FEATURES_FILE, 'w') as f:
        json.dump(features, f)

# 加载人脸数据
def load_faces():
    if os.path.exists(FACES_FILE):
        with open(FACES_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    return []

# 保存人脸数据
def save_faces(faces):
    with open(FACES_FILE, 'w', encoding='utf-8') as f:
        json.dump(faces, f, ensure_ascii=False, indent=2)

# 加载人名数据
def load_face_names():
    if os.path.exists(FACE_NAMES_FILE):
        with open(FACE_NAMES_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    return []

# 保存人名数据
def save_face_names(names):
    with open(FACE_NAMES_FILE, 'w', encoding='utf-8') as f:
        json.dump(names, f, ensure_ascii=False, indent=2)

# 获取人名列表
@app.get("/api/face-names")
async def get_face_names(username: str = Depends(verify_token)):
    return load_face_names()

# 添加新的人名
@app.post("/api/face-names")
async def create_face_name(
    data: dict,
    username: str = Depends(verify_token)
):
    try:
        names = load_face_names()
        new_name = {
            "id": len(names) + 1,
            "name": data["name"],
            "created_at": datetime.now().isoformat(),
            "created_by": username
        }
        names.append(new_name)
        save_face_names(names)
        
        # 创建对应的文件夹
        create_face_folder(str(new_name["id"]))
        
        return new_name
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 更新人脸的人名
@app.put("/api/faces/{face_id}/name")
async def update_face_name(
    face_id: str,
    data: dict,
    username: str = Depends(verify_token)
):
    faces = load_faces()
    face = next((f for f in faces if f["id"] == face_id), None)
    
    if not face:
        raise HTTPException(status_code=404, detail="人脸不存在")
    
    face["name_id"] = data.get("name_id")
    face["name"] = data.get("name")
    save_faces(faces)
    
    return {"message": "更新成功"}

# 初始化人名数据
def init_face_names():
    if not os.path.exists(FACE_NAMES_FILE):
        with open(FACE_NAMES_FILE, 'w', encoding='utf-8') as f:
            json.dump([], f, ensure_ascii=False)

# 删除人名
@app.delete("/api/face-names/{name_id}")
async def delete_face_name(
    name_id: int,
    username: str = Depends(verify_token)
):
    try:
        names = load_face_names()
        names = [n for n in names if n["id"] != name_id]
        save_face_names(names)
        
        # 删除对应的文件夹及其内容
        folder_path = os.path.join(FACE_DB_ROOT, str(name_id))
        if os.path.exists(folder_path):
            shutil.rmtree(folder_path)
        
        return {"message": "删除成功"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 上传训练图片
@app.post("/api/face-db/{name_id}")
async def upload_training_face(
    name_id: str,  # 先接收为字符串
    file: UploadFile = File(...),
    username: str = Depends(verify_token)
):
    try:
        # 获取人名信息
        names = load_face_names()
        # 将字符串转换为整数进行比较
        name_id_int = int(name_id)
        name_info = next((n for n in names if n["id"] == name_id_int), None)
        if not name_info:
            raise HTTPException(status_code=404, detail="人名不存在")
        
        # 保存到对应人名的文件夹
        folder_path = os.path.join(FACE_DB_ROOT, name_id)  # 使用原始字符串
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)
        
        # 生成文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        random_str = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=8))
        filename = f"{timestamp}_{random_str}{os.path.splitext(file.filename)[1]}"
        file_path = os.path.join(folder_path, filename)
        
        # 保存文件
        with open(file_path, "wb") as f:
            shutil.copyfileobj(file.file, f)
        
        # 提取并保存人脸特征
        face_features = get_face_encodings(file_path)
        if not face_features:
            os.remove(file_path)
            raise HTTPException(status_code=400, detail="未检测到人脸")
        
        return {"message": "上传成功", "features": face_features}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 识别人脸
@app.post("/api/recognize")
async def recognize_face(
    file: UploadFile = File(...),
    username: str = Depends(verify_token)
):
    try:
        # 保存上传的图片
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        random_str = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=8))
        filename = f"{timestamp}_{random_str}{os.path.splitext(file.filename)[1]}"
        file_path = os.path.join(UPLOAD_DIRS['face'], filename)
        
        with open(file_path, "wb") as f:
            shutil.copyfileobj(file.file, f)
        
        # 检测人脸并提取特征
        face_features = get_face_encodings(file_path)
        if not face_features:
            raise HTTPException(status_code=400, detail="未检测到人脸")
        
        # 与数据库中的人脸比对
        results = []
        for face in face_features:
            best_match = None
            best_score = float('inf')
            matched_name = None
            matched_name_info = None
            
            # 获取所有人名信息
            names = load_face_names()
            for name_info in names:
                folder_path = os.path.join(FACE_DB_ROOT, str(name_info["id"]))
                if os.path.isdir(folder_path):
                    # 遍历文件夹中的所有图片
                    for img_file in os.listdir(folder_path):
                        if img_file.lower().endswith(('.png', '.jpg', '.jpeg')):
                            img_path = os.path.join(folder_path, img_file)
                            try:
                                db_features = get_face_encodings(img_path)
                                
                                for db_face in db_features:
                                    score = sum((a - b) ** 2 for a, b in 
                                              zip(face['feature'], db_face['feature']))
                                    if score < best_score:
                                        best_score = score
                                        best_match = db_face
                                        matched_name = name_info["name"]
                                        matched_name_info = name_info
                            except Exception as e:
                                print(f"处理图片失败: {img_path}, 错误: {str(e)}")
                                continue
            
            results.append({
                'bbox': face['bbox'],
                'name': matched_name if best_score < 1000 else "未知",  # 调整阈值
                'name_id': matched_name_info["id"] if matched_name_info else None,
                'confidence': 1 / (1 + best_score) if best_score < float('inf') else 0,
                'attributes': {
                    'age': '未知',
                    'gender': '未知',
                    'emotion': '未知'
                }
            })
        
        return {
            'faces': results,
            'original_image': f"faces/{filename}",
        }
    except Exception as e:
        print(f"识别失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 获取人名下的所有图片
@app.get("/api/face-db/{name_id}/images")
async def get_name_images(
    name_id: str,
    username: str = Depends(verify_token)
):
    try:
        # 检查人名是否存在
        names = load_face_names()
        name_id_int = int(name_id)
        name_info = next((n for n in names if n["id"] == name_id_int), None)
        if not name_info:
            raise HTTPException(status_code=404, detail="人名不存在")
        
        # 获取文件夹中的所有图片
        folder_path = os.path.join(FACE_DB_ROOT, name_id)
        if not os.path.exists(folder_path):
            return []
        
        images = []
        for filename in os.listdir(folder_path):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                images.append({
                    'filename': filename,
                    'url': f"face_db/{name_id}/{filename}",
                    'created_at': datetime.fromtimestamp(
                        os.path.getctime(os.path.join(folder_path, filename))
                    ).isoformat()
                })
        
        return images
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 删除训练图片
@app.delete("/api/face-db/{name_id}/images/{filename}")
async def delete_training_image(
    name_id: str,
    filename: str,
    username: str = Depends(verify_token)
):
    try:
        # 检查人名是否存在
        names = load_face_names()
        name_id_int = int(name_id)
        name_info = next((n for n in names if n["id"] == name_id_int), None)
        if not name_info:
            raise HTTPException(status_code=404, detail="人名不存在")
        
        # 删除图片文件
        file_path = os.path.join(FACE_DB_ROOT, name_id, filename)
        print(f"尝试删除文件: {file_path}")  # 添加调试日志
        if os.path.exists(file_path):
            os.remove(file_path)
            return {"message": "删除成功"}
        else:
            print(f"文件不存在: {file_path}")  # 添加调试日志
            raise HTTPException(status_code=404, detail="图片不存在")
    except Exception as e:
        print(f"删除失败: {str(e)}")  # 添加调试日志
        raise HTTPException(status_code=500, detail=str(e))

def detect_objects_in_image(image_path):
    """模拟物体检测"""
    try:
        # 读取图片
        image = cv2.imread(image_path)
        height, width = image.shape[:2]
        
        # 生成随机的检测结果
        objects = []
        # 随机生成2-5个物体
        num_objects = random.randint(2, 5)
        for i in range(num_objects):
            # 随机选择类别和置信度
            category = random.choice(MOCK_CLASSES)
            confidence = random.uniform(0.6, 0.99)
            
            # 随机生成边界框（确保在图片范围内）
            w = random.randint(50, 150)
            h = random.randint(50, 150)
            x = random.randint(0, width - w)
            y = random.randint(0, height - h)
            
            objects.append({
                "id": i + 1,
                "category": category,
                "confidence": confidence,
                "bbox": {
                    "x": float(x/width * 100),
                    "y": float(y/height * 100),
                    "width": float(w/width * 100),
                    "height": float(h/height * 100)
                }
            })
        
        return objects
    except Exception as e:
        print(f"物体检测失败: {str(e)}")
        return []

# 物体检测路由
@app.post("/detect_objects")
async def detect_objects(
    file: UploadFile = File(...),
    username: str = Depends(verify_token)
):
    try:
        # 保存上传的图片
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        random_str = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=8))
        filename = f"{timestamp}_{random_str}{os.path.splitext(file.filename)[1]}"
        file_path = os.path.join(UPLOAD_DIRS['objects'], filename)
        
        with open(file_path, "wb") as f:
            shutil.copyfileobj(file.file, f)
        
        # 检测物体
        objects = detect_objects_in_image(file_path)
        
        # 在图片上标注检测结果
        image = cv2.imread(file_path)
        for obj in objects:
            bbox = obj['bbox']
            x = int(bbox['x'] * image.shape[1] / 100)
            y = int(bbox['y'] * image.shape[0] / 100)
            w = int(bbox['width'] * image.shape[1] / 100)
            h = int(bbox['height'] * image.shape[0] / 100)
            
            cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
            label = f"{obj['category']} ({obj['confidence']:.2%})"
            cv2.putText(image, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # 保存标注后的图片
        annotated_filename = f"annotated_{filename}"
        annotated_path = os.path.join(UPLOAD_DIRS['objects'], annotated_filename)
        cv2.imwrite(annotated_path, image)
        
        return {
            "original_image": f"objects/{filename}",
            "annotated_image": f"objects/{annotated_filename}",
            "objects": objects
        }
    except Exception as e:
        print(f"处理失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 导出数据集
@app.get("/api/export")
async def export_dataset(
    type: str = Query(default="category", description="数据集类型"),
    username: str = Depends(verify_token)
):
    try:
        # 创建临时ZIP文件
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            # 读取图片数据和分类信息
            images = load_images(type)
            categories = load_categories() if type == "category" else []
            print(f"正在导出数据集，共 {len(images)} 个文件")
            
            # 创建数据集目录结构
            dataset_structure = {
                "categories": {},  # 按分类组织的图片
                "uncategorized": [],  # 未分类的图片
                "metadata": []  # 元数据
            }
            
            # 添加图片文件
            for image in images:
                # 获取文件名（不含路径）
                filename = os.path.basename(image["filename"])
                
                # 原始图片路径
                orig_path = os.path.join(UPLOAD_DIRS[type], filename)
                if os.path.exists(orig_path):
                    # 根据分类添加到相应目录
                    category_id = image.get("category_id")
                    if category_id:
                        # 获取分类信息
                        category = next((c for c in categories if c["id"] == category_id), None)
                        if category:
                            category_name = category["name"]
                            # 确保分类名称可以用作文件夹名
                            safe_category_name = "".join(c for c in category_name if c.isalnum() or c in (' ', '_', '-')).strip()
                            # 添加到对应分类目录
                            zip_path = f"dataset/{safe_category_name}/{filename}"
                            zip_file.write(orig_path, zip_path)
                            # 记录到结构中
                            if safe_category_name not in dataset_structure["categories"]:
                                dataset_structure["categories"][safe_category_name] = []
                            dataset_structure["categories"][safe_category_name].append(filename)
                    else:
                        # 未分类的图片
                        zip_path = f"dataset/未分类/{filename}"
                        zip_file.write(orig_path, zip_path)
                        dataset_structure["uncategorized"].append(filename)
                      
                    # 添加元数据
                    metadata = {
                        "filename": filename,
                        "category_id": image.get("category_id"),
                        "category_name": next((c["name"] for c in categories if c["id"] == image.get("category_id")), "未分类"),
                        "uploaded_by": image.get("uploaded_by"),
                        "uploaded_at": image.get("uploaded_at")
                    }
                    dataset_structure["metadata"].append(metadata)
            
            # 添加分类信息
            if type == "category":
                categories = load_categories()
                zip_file.writestr('dataset/categories.json', 
                                json.dumps(categories, ensure_ascii=False, indent=2))
            
            # 添加数据集信息文件
            dataset_info = {
                "total_images": len(images),
                "categories": {
                    name: len(files) for name, files in dataset_structure["categories"].items()
                },
                "uncategorized": len(dataset_structure["uncategorized"]),
                "export_time": datetime.now().isoformat(),
                "exported_by": username,
                "type": type
            }
            zip_file.writestr('dataset/dataset_info.json', 
                            json.dumps(dataset_info, ensure_ascii=False, indent=2))
            
            # 添加元数据文件
            zip_file.writestr('dataset/metadata.json',
                            json.dumps(dataset_structure["metadata"], 
                                     ensure_ascii=False, indent=2))
        
        zip_buffer.seek(0)
        
        # 生成下载文件名
        filename = f"dataset_{type}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
        print(f"数据集导出完成: {filename}")
        
        return StreamingResponse(
            iter([zip_buffer.getvalue()]),
            media_type="application/zip",
            headers={
                "Content-Disposition": f"attachment; filename={filename}"
            }
        )
    except Exception as e:
        print(f"导出数据集失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

def detect_persons_in_image(image_path):
    """行人检测函数"""
    try:
        # 读取图片
        image = cv2.imread(image_path)
        height, width = image.shape[:2]
        
        # 调整图片大小以提高检测效果
        scale = min(800 / width, 800 / height)
        if scale < 1:
            width = int(width * scale)
            height = int(height * scale)
            image = cv2.resize(image, (width, height))
        
        # 使用 HOG 进行行人检测
        boxes, weights = hog.detectMultiScale(
            image,
            winStride=(8, 8),
            padding=(8, 8),
            scale=1.05,
            useMeanshiftGrouping=False
        )
        
        # 非极大值抑制，合并重叠的框
        boxes = np.array([[x, y, x + w, y + h] for (x, y, w, h) in boxes])
        if len(boxes) > 0:
            pick = non_max_suppression(boxes, probs=weights, overlapThresh=0.65)
        else:
            pick = []
        
        # 模拟行人检测结果
        persons = []
        for i, (xA, yA, xB, yB) in enumerate(pick):
            w = xB - xA
            h = yB - yA
            confidence = float(weights[i])
            
            # 提取行人区域进行进一步分析
            roi = image[yA:yB, xA:xB]
            
            # 根据姿态估计判断姿势
            aspect_ratio = float(h) / w
            if aspect_ratio > 2.5:
                pose = '站立'
            elif aspect_ratio > 2.0:
                pose = '行走'
            else:
                pose = '其他'
            
            # 根据运动检测判断动作
            if len(roi) > 0:
                gray_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
                blur = cv2.GaussianBlur(gray_roi, (21, 21), 0)
                _, thresh = cv2.threshold(blur, 20, 255, cv2.THRESH_BINARY)
                if np.mean(thresh) > 127:
                    action = '移动'
                else:
                    action = '静止'
            else:
                action = '未知'
            
            # 分析服装颜色
            if len(roi) > 0:
                mean_color = cv2.mean(roi)[:3]
                if mean_color[0] > 200 and mean_color[1] > 200 and mean_color[2] > 200:
                    clothing = '浅色服装'
                elif mean_color[0] < 100 and mean_color[1] < 100 and mean_color[2] < 100:
                    clothing = '深色服装'
                else:
                    clothing = '中等色调'
            else:
                clothing = '未知'
            
            person = {
                "id": i + 1,
                "confidence": float(confidence),
                "bbox": {
                    "x": float(xA/width * 100),
                    "y": float(yA/height * 100),
                    "width": float(w/width * 100),
                    "height": float(h/height * 100)
                },
                "pose": pose,
                "action": action,
                "clothing": clothing
            }
            persons.append(person)
        
        # 在图片上标注行人
        annotated_image = image.copy()
        for person in persons:
            bbox = person['bbox']
            xA = int(bbox['x'] * width / 100)
            yA = int(bbox['y'] * height / 100)
            w = int(bbox['width'] * width / 100)
            h = int(bbox['height'] * height / 100)
            
            # 绘制边界框
            color = (
                int(255 * (1 - person['confidence'])),
                int(255 * person['confidence']),
                0
            )
            cv2.rectangle(annotated_image, (xA, yA), (xA + w, yA + h), color, 2)
            
            # 添加标签
            label = f"Person {person['id']} ({person['confidence']:.0%})"
            label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
            cv2.rectangle(annotated_image,
                        (xA, yA - 20),
                        (xA + label_size[0], yA),
                        (0, 0, 0),
                        -1)
            cv2.putText(annotated_image, label, (xA, yA - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        
        return persons, annotated_image
        
    except Exception as e:
        print(f"行人检测失败: {str(e)}")
        traceback.print_exc()
        return [], None

def non_max_suppression(boxes, probs=None, overlapThresh=0.3):
    """非极大值抑制，用于合并重叠的检测框"""
    if len(boxes) == 0:
        return []

    if boxes.dtype.kind == "i":
        boxes = boxes.astype("float")

    pick = []

    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = y2

    if probs is not None:
        idxs = probs

    idxs = np.argsort(idxs)

    while len(idxs) > 0:
        last = len(idxs) - 1
        i = idxs[last]
        pick.append(i)

        xx1 = np.maximum(x1[i], x1[idxs[:last]])
        yy1 = np.maximum(y1[i], y1[idxs[:last]])
        xx2 = np.minimum(x2[i], x2[idxs[:last]])
        yy2 = np.minimum(y2[i], y2[idxs[:last]])

        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)

        overlap = (w * h) / area[idxs[:last]]

        idxs = np.delete(idxs, np.concatenate(([last],
            np.where(overlap > overlapThresh)[0])))

    return boxes[pick].astype("int")

# 行人检测上传接口
@app.post("/api/persons/upload")
async def upload_person_image(
    file: UploadFile = File(...),
    username: str = Depends(verify_token)
):
    try:
        # 确保目录存在
        if not os.path.exists(UPLOAD_DIRS['persons']):
            os.makedirs(UPLOAD_DIRS['persons'])

        # 验证文件类型
        if not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="只能上传图片文件")
        
        # 生成文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        random_str = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=8))
        filename = f"{timestamp}_{random_str}{os.path.splitext(file.filename)[1]}"
        
        # 保存原始图片
        file_path = os.path.join(UPLOAD_DIRS['persons'], filename)
        with open(file_path, "wb") as f:
            shutil.copyfileobj(file.file, f)
        
        # 进行行人检测
        persons, annotated_image = detect_persons_in_image(file_path)
        if persons is None or annotated_image is None:
            raise HTTPException(status_code=500, detail="行人检测失败")
        
        # 保存标注后的图片
        annotated_filename = f"annotated_{filename}"
        annotated_path = os.path.join(UPLOAD_DIRS['persons'], annotated_filename)
        cv2.imwrite(annotated_path, annotated_image, [cv2.IMWRITE_JPEG_QUALITY, 95])
        
        # 读取现有图片数据
        images = load_images('persons')
        if images is None:
            images = []
        
        # 创建新的图片记录
        new_image = {
            "id": len(images) + 1,
            "filename": f"persons/{filename}",
            "annotated_image": f"persons/{annotated_filename}",
            "persons": persons,
            "uploaded_by": username,
            "uploaded_at": datetime.now().isoformat()
        }
        
        # 保存图片数据
        images.append(new_image)
        save_images(images, 'persons')
        
        return new_image
        
    except Exception as e:
        print(f"上传失败: {str(e)}")
        traceback.print_exc()
        # 清理可能的临时文件
        if 'file_path' in locals() and os.path.exists(file_path):
            try:
                os.remove(file_path)
            except:
                pass
        raise HTTPException(status_code=500, detail=str(e))

# 获取行人检测图片列表
@app.get("/api/persons")
async def get_person_images(username: str = Depends(verify_token)):
    try:
        # 读取行人检测图片数据
        images = load_images('persons')
        if images is None:
            images = []
        
        # 确保文件路径正确
        for image in images:
            if "filename" in image and not image["filename"].startswith("persons"):
                image["filename"] = f"persons/{os.path.basename(image['filename'])}"
            if "annotated_image" in image and not image["annotated_image"].startswith("persons"):
                image["annotated_image"] = f"persons/{os.path.basename(image['annotated_image'])}"
        
        return images
        
    except Exception as e:
        print(f"获取行人检测图片列表失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 删除行人检测图片
@app.delete("/api/persons/{image_id}")
async def delete_person_image(
    image_id: int,
    username: str = Depends(verify_token)
):
    try:
        # 读取图片数据
        images = load_images('persons')
        if not images:
            raise HTTPException(status_code=404, detail="没有找到图片")
        
        # 查找要删除的图片
        image = next((img for img in images if img["id"] == image_id), None)
        if not image:
            raise HTTPException(status_code=404, detail="图片不存在")
        
        # 删除文件
        if "filename" in image:
            filename = os.path.basename(image["filename"])
            file_path = os.path.join(UPLOAD_DIRS['persons'], filename)
            if os.path.exists(file_path):
                os.remove(file_path)
        
        # 删除标注图片
        if "annotated_image" in image:
            annotated_filename = os.path.basename(image["annotated_image"])
            annotated_path = os.path.join(UPLOAD_DIRS['persons'], annotated_filename)
            if os.path.exists(annotated_path):
                os.remove(annotated_path)
        
        # 从列表中移除
        images = [img for img in images if img["id"] != image_id]
        save_images(images, 'persons')
        
        return {"message": "图片删除成功"}
        
    except Exception as e:
        print(f"删除图片失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# 导出行人检测数据集
@app.get("/api/persons/export")
async def export_person_dataset(username: str = Depends(verify_token)):
    try:
        # 创建临时ZIP文件
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            # 读取图片数据
            images = load_images('persons')
            if not images:
                raise HTTPException(status_code=404, detail="没有可导出的数据")
            
            print(f"正在导出行人检测数据集，共 {len(images)} 个文件")
            
            # 创建数据集目录结构
            dataset = {
                "info": {
                    "description": "行人检测数据集",
                    "version": "1.0",
                    "year": datetime.now().year,
                    "contributor": username,
                    "date_created": datetime.now().isoformat()
                },
                "images": [],
                "annotations": []
            }
            
            annotation_id = 1
            
            # 处理每张图片
            for image in images:
                # 获取文件名
                filename = os.path.basename(image["filename"])
                
                # 添加图片信息
                image_info = {
                    "id": image["id"],
                    "file_name": filename,
                    "uploaded_by": image["uploaded_by"],
                    "uploaded_at": image["uploaded_at"]
                }
                dataset["images"].append(image_info)
                
                # 添加标注信息
                if "persons" in image:
                    for person in image["persons"]:
                        annotation = {
                            "id": annotation_id,
                            "image_id": image["id"],
                            "category_id": 1,  # 1 表示行人
                            "bbox": [
                                person["bbox"]["x"],
                                person["bbox"]["y"],
                                person["bbox"]["width"],
                                person["bbox"]["height"]
                            ],
                            "confidence": person["confidence"],
                            "attributes": {
                                "pose": person["pose"],
                                "action": person["action"],
                                "clothing": person["clothing"]
                            }
                        }
                        dataset["annotations"].append(annotation)
                        annotation_id += 1
                
                # 复制原始图片
                orig_path = os.path.join(UPLOAD_DIRS["persons"], filename)
                if os.path.exists(orig_path):
                    zip_file.write(orig_path, f"images/{filename}")
                
                # 复制标注图片
                if "annotated_image" in image:
                    annotated_filename = os.path.basename(image["annotated_image"])
                    annotated_path = os.path.join(UPLOAD_DIRS["persons"], annotated_filename)
                    if os.path.exists(annotated_path):
                        zip_file.write(annotated_path, f"visualizations/{annotated_filename}")
            
            # 添加数据集描述文件
            zip_file.writestr('dataset.json', 
                            json.dumps(dataset, ensure_ascii=False, indent=2))
            
            # 添加README文件
            readme_content = f"""# 行人检测数据集

## 数据集信息
- 描述：{dataset['info']['description']}
- 版本：{dataset['info']['version']}
- 创建时间：{dataset['info']['date_created']}
- 贡献者：{dataset['info']['contributor']}

## 统计信息
- 图片总数：{len(dataset['images'])}
- 标注总数：{len(dataset['annotations'])}

## 目录结构
- images/: 原始图片
- visualizations/: 标注可视化结果
- dataset.json: 数据集描述文件
"""
            zip_file.writestr('README.md', readme_content)
        
        zip_buffer.seek(0)
        
        # 生成下载文件名
        filename = f"person_detection_dataset_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
        print(f"数据集导出完成: {filename}")
        
        return StreamingResponse(
            iter([zip_buffer.getvalue()]),
            media_type="application/zip",
            headers={
                "Content-Disposition": f"attachment; filename={filename}"
            }
        )
        
    except Exception as e:
        print(f"导出数据集失败: {str(e)}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000) 